{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimize and Deploy AI Models with Pruna and Hugging Face\n",
    "\n",
    "Objective: Build a complete tutorial demonstrating how to optimize the [Efficient-Large-Model/Sana_600M_512px_diffusers](https://huggingface.co/Efficient-Large-Model/Sana_600M_512px_diffusers) diffusion model using Pruna and deploy it seamlessly to the Hugging Face Hub.\n",
    "\n",
    "Model: [Efficient-Large-Model/Sana_600M_512px_diffusers](https://huggingface.co/Efficient-Large-Model/Sana_600M_512px_diffusers)\n",
    "\n",
    "Dataset: [data-is-better-together/open-image-preferences-v1-binarized](https://huggingface.co/datasets/data-is-better-together/open-image-preferences-v1-binarized)\n",
    "\n",
    "To follow along, ensure that you have the Pruna SDK installed along with all required third-party libraries. Running this tutorial in a clean virtual environment is recommended for a smooth setup."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "outputs": [],
   "source": [
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/PrunaAI/pruna/blob/v|version|/docs/tutorials/deploying_sana_tutorial.ipynb\">\n",
    "    <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
    "</a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install pruna"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install datasets huggingface_hub gradio"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You will need to login on the Hugging Face Hub for using the model weights. We also need to select the best available device for executing the notebook. Run the cells below to do the same."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import notebook_login\n",
    "\n",
    "notebook_login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"mps\" if torch.backends.mps.is_available() else \"cpu\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Smash Configuration\n",
    "\n",
    "To optimize the model effectively, we first define the configuration methods that enhance its performance. For detailed options and parameter explanations, refer to the [SmashConfig guide](https://docs.pruna.ai/en/stable/docs_pruna/user_manual/configure.html).\n",
    "\n",
    "In this tutorial, we will:\n",
    "\n",
    "* Select **torchao** to reduce memory usage during inference.\n",
    "*  Save the model to disk and optionally upload the smashed model to the Hugging Face Hub for easy access."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Multiple distributions found for package optimum. Picked distribution: optimum-quanto\n",
      "/teamspace/studios/this_studio/.venv/lib/python3.11/site-packages/torch/cuda/__init__.py:61: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.\n",
      "  import pynvml  # type: ignore[import]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9714bd1c0ad4487586b84bb63ddcfb47",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model_index.json:   0%|          | 0.00/401 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "89e695e1dca64729b16bd016d8a45a93",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Fetching 14 files:   0%|          | 0/14 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "41e6d1e4bbf74deeb88787fd0013ad13",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "text_encoder/model.fp16-00001-of-00002.s(…):   0%|          | 0.00/4.99G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f7c8b274ec3c4714b1d774631368343e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "scheduler_config.json:   0%|          | 0.00/840 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "679ea9ac5f4d4c1aaf2146df0754ea85",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "config.json:   0%|          | 0.00/875 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bd9313dea5df44618fc71fb3bed2c428",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer/tokenizer.model:   0%|          | 0.00/4.24M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b24b533a028442fe860a8ccadb8cfe0a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer/tokenizer.json:   0%|          | 0.00/34.4M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b25f8769acad47f3ba137572735244cf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "text_encoder/model.fp16-00002-of-00002.s(…):   0%|          | 0.00/241M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e6d990d1b12f427392b33222df67a3fc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model.safetensors.index.fp16.json: 0.00B [00:00, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "91e6f92a6b8445c3a38a34333de53376",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "special_tokens_map.json:   0%|          | 0.00/636 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ae40e43578cd4d84969e1b98cd6030ad",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "transformer/diffusion_pytorch_model.fp16(…):   0%|          | 0.00/1.18G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4d210e72a74544618e1266d6f4744bb1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer_config.json: 0.00B [00:00, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "71be7a256f1b40508ab3be54a5fb71d2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "config.json:   0%|          | 0.00/495 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bc2424caa4a44d77921ca782e700203f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "config.json: 0.00B [00:00, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6d1290aa460e4184a08394cf06752dc9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "vae/diffusion_pytorch_model.fp16.safeten(…):   0%|          | 0.00/1.25G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d25c7a12fbc44303bda12f8e3219595a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading pipeline components...:   0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2bb0797430fe49f48ac5b125ba6591e9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO - Using best available device: 'cuda'\n",
      "INFO - Starting quantizer torchao...\n",
      "INFO - quantizer torchao was applied successfully.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ Smashed Sana model uploaded successfully to Hugging Face Hub.\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from diffusers import SanaPipeline\n",
    "\n",
    "from pruna import PrunaModel, SmashConfig, smash\n",
    "\n",
    "# 1. Define the model ID\n",
    "model_id = \"Efficient-Large-Model/Sana_600M_512px_diffusers\"\n",
    "\n",
    "# 2. Load the pre-trained model\n",
    "pipe = SanaPipeline.from_pretrained(model_id, variant=\"fp16\", torch_dtype=torch.float16)\n",
    "pipe = pipe.to(device)\n",
    "\n",
    "# 3. Configure Pruna smash\n",
    "smash_config = SmashConfig([\"torchao\"])\n",
    "\n",
    "# 4. Smash (optimize) the model\n",
    "smashed_pipe = smash(model=pipe, smash_config=smash_config)\n",
    "\n",
    "# 5. Save the model to disk\n",
    "smashed_pipe.save_pretrained(\"Sana_600M_512px_diffusers-smashed\")\n",
    "\n",
    "print(\"✅ Smashed Sana model uploaded successfully to Hugging Face Hub.\")\n",
    "\n",
    "# 6. Push the smashed pipeline to Hugging Face Hub using save_to_hub\n",
    "# smashed_pipe.save_to_hub(\"AINovice2005/Sana_600M_512px_diffusers-smashed\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## **2. Load and Collate Dataset**\n",
    "\n",
    "In this step, we will load the dataset required for optimizing and evaluating the model. This dataset will provide the input data needed to assess the model’s performance after applying optimization techniques such as quantization.\n",
    "\n",
    "We will use the [`data-is-better-together/open-image-preferences-v1-binarized`](https://huggingface.co/datasets/data-is-better-together/open-image-preferences-v1-binarized) dataset, which contains binarized user image preferences and prompts for image generation tasks. Correctly loading and collating the dataset ensures that the input is properly prepared, enabling smooth evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO - Loaded only training, splitting train 80/10/10 into train, validation and test...\n",
      "INFO - Testing compatibility with image_generation_collate...\n"
     ]
    }
   ],
   "source": [
    "from datasets import Image, load_dataset\n",
    "\n",
    "from pruna.data.pruna_datamodule import PrunaDataModule\n",
    "from pruna.data.utils import split_train_into_train_val_test\n",
    "\n",
    "# Load dataset\n",
    "dataset = load_dataset(\"data-is-better-together/open-image-preferences-v1-binarized\")[\"train\"]\n",
    "\n",
    "dataset = dataset.rename_column(\"chosen\", \"image\")\n",
    "dataset = dataset.rename_column(\"prompt\", \"text\")\n",
    "\n",
    "dataset = dataset.cast_column(\"image\", Image())\n",
    "\n",
    "# Split train into train/val/test\n",
    "train_ds, val_ds, test_ds = split_train_into_train_val_test(dataset, seed=42)\n",
    "\n",
    "# Initialize PrunaDataModule\n",
    "datamodule = PrunaDataModule.from_datasets(\n",
    "    datasets=(train_ds, val_ds, test_ds),\n",
    "    collate_fn=\"image_generation_collate\",\n",
    "    collate_fn_args={\"img_size\": 512},\n",
    ")\n",
    "\n",
    "# Limit datasets to 5 samples each for quick testing\n",
    "datamodule.limit_datasets(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Evaluate the Model\n",
    "\n",
    "Now that the model and dataset are set up, we can proceed to evaluate the model using the **Pruna Evaluation Agent**. This evaluation helps us measure the model’s current performance before optimization, providing a baseline for comparison. It assesses how well the model performs on the given dataset and generates relevant metrics that will guide us in understanding the impact of our optimization configurations later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO - Using best available device: 'cuda'\n",
      "WARNING - Argument cache_dir not found in config file. Skipping...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4bbd538e86fc4f6dac612f15fec1d3c1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading pipeline components...:   0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a7731367b3de485aa99c194b72f79e9b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING - Model and SmashConfig have different devices. Model: cuda, SmashConfig: cuda:0. Casting model to cuda:0.If this is not desired, please use SmashConfig(device='cuda').\n",
      "INFO - Starting quantizer torchao...\n",
      "INFO - quantizer torchao was applied successfully.\n",
      "INFO - Using best available device: 'cuda'\n",
      "INFO - Using best available device: 'cuda'\n",
      "INFO - Using provided list of metric instances.\n",
      "INFO - Using best available device: 'cuda'\n",
      "INFO - Evaluating a smashed model.\n",
      "INFO - Detected diffusers model. Using DiffuserHandler with fixed seed.\n",
      "- The first element of the batch is passed as input.\n",
      "- The generated outputs are expected to have .images attribute.\n",
      "Pipelines loaded with `dtype=torch.float16` cannot run with `cpu` device. It is not recommended to move them to `cpu` as running them will fail. Please make sure to use an accelerator to run the pipeline in inference, due to the lack of support for`float16` operations on this device in PyTorch. Please, remove the `torch_dtype=torch.float16` argument, or use another device for inference.\n",
      "Pipelines loaded with `dtype=torch.float16` cannot run with `cpu` device. It is not recommended to move them to `cpu` as running them will fail. Please make sure to use an accelerator to run the pipeline in inference, due to the lack of support for`float16` operations on this device in PyTorch. Please, remove the `torch_dtype=torch.float16` argument, or use another device for inference.\n",
      "Pipelines loaded with `dtype=torch.float16` cannot run with `cpu` device. It is not recommended to move them to `cpu` as running them will fail. Please make sure to use an accelerator to run the pipeline in inference, due to the lack of support for`float16` operations on this device in PyTorch. Please, remove the `torch_dtype=torch.float16` argument, or use another device for inference.\n",
      "INFO - Evaluating stateful metrics.\n",
      "INFO - Evaluating isolated inference metrics.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "68274baf8a254d9e9cbb2d14c460af54",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/20 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6f950e45731841fbacc61d9cfb6ca64e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/20 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Pipelines loaded with `dtype=torch.float16` cannot run with `cpu` device. It is not recommended to move them to `cpu` as running them will fail. Please make sure to use an accelerator to run the pipeline in inference, due to the lack of support for`float16` operations on this device in PyTorch. Please, remove the `torch_dtype=torch.float16` argument, or use another device for inference.\n",
      "Pipelines loaded with `dtype=torch.float16` cannot run with `cpu` device. It is not recommended to move them to `cpu` as running them will fail. Please make sure to use an accelerator to run the pipeline in inference, due to the lack of support for`float16` operations on this device in PyTorch. Please, remove the `torch_dtype=torch.float16` argument, or use another device for inference.\n",
      "Pipelines loaded with `dtype=torch.float16` cannot run with `cpu` device. It is not recommended to move them to `cpu` as running them will fail. Please make sure to use an accelerator to run the pipeline in inference, due to the lack of support for`float16` operations on this device in PyTorch. Please, remove the `torch_dtype=torch.float16` argument, or use another device for inference.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[MetricResult(name='total_time', params={'n_iterations': 1, 'n_warmup_iterations': 1, 'device': 'cuda', 'timing_type': 'sync', 'batch_size': 1}, result=12479.54296875), MetricResult(name='latency', params={'n_iterations': 1, 'n_warmup_iterations': 1, 'device': 'cuda', 'timing_type': 'sync', 'batch_size': 1}, result=12479.54296875)]\n"
     ]
    }
   ],
   "source": [
    "# Import required modules from Pruna\n",
    "from pruna.data.pruna_datamodule import PrunaDataModule\n",
    "from pruna.engine.utils import move_to_device\n",
    "from pruna.evaluation.evaluation_agent import EvaluationAgent\n",
    "from pruna.evaluation.metrics import (\n",
    "    LatencyMetric,\n",
    "    TotalTimeMetric,\n",
    ")\n",
    "from pruna.evaluation.task import Task\n",
    "\n",
    "# Load the smashed (optimized) model pipeline from disk\n",
    "smashed_pipe = PrunaModel.from_pretrained(\"Sana_600M_512px_diffusers-smashed\")\n",
    "\n",
    "# Define evaluation metrics (example: total time and latency)\n",
    "metrics = [\n",
    "    TotalTimeMetric(n_iterations=1, n_warmup_iterations=1),\n",
    "    LatencyMetric(n_iterations=1, n_warmup_iterations=1),\n",
    "]\n",
    "\n",
    "# Define the evaluation task with metrics and datamodule\n",
    "# (Ensure `datamodule` and `device` are defined before this script runs)\n",
    "task = Task(metrics, datamodule=datamodule, device=device)\n",
    "\n",
    "# Initialize the evaluation agent\n",
    "eval_agent = EvaluationAgent(task)\n",
    "\n",
    "# Move smashed model to evaluation device (GPU or CPU)\n",
    "move_to_device(smashed_pipe, device)\n",
    "\n",
    "# Evaluate the smashed model pipeline using the evaluation agent\n",
    "smashed_model_results = eval_agent.evaluate(smashed_pipe)\n",
    "\n",
    "# Print results for verification\n",
    "print(smashed_model_results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Gradio Demo\n",
    "\n",
    "Once the model has been optimized, we can deploy the smashed model using **Gradio** to create an interactive demo. This allows anyone to test the model’s capabilities directly in their browser.\n",
    "\n",
    "In this section, we will:\n",
    "\n",
    "* Show how to deploy the optimized model on the Hugging Face Hub with a Gradio demo\n",
    "* Discuss considerations such as **handling queuing**, especially if multiple users access the demo simultaneously\n",
    "* Highlight best practices for integrating Gradio demos in your Hugging Face Space to ensure a smooth and responsive user experience\n",
    "\n",
    "Creating a Gradio demo not only showcases your optimized model effectively but also enables easy sharing and real-world testing by the community."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO - Using best available device: 'cuda'\n",
      "WARNING - Argument cache_dir not found in config file. Skipping...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "229f7cd246da4edfa7f9c9cd47d98631",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading pipeline components...:   0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "96cd717cb19c4275a4a78bc6a28b69e6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING - Model and SmashConfig have different devices. Model: cuda, SmashConfig: cuda:0. Casting model to cuda:0.If this is not desired, please use SmashConfig(device='cuda').\n",
      "INFO - Starting quantizer torchao...\n",
      "INFO - quantizer torchao was applied successfully.\n",
      "/teamspace/studios/this_studio/.venv/lib/python3.11/site-packages/gradio/interface.py:415: UserWarning: The `allow_flagging` parameter in `Interface` is deprecated. Use `flagging_mode` instead.\n",
      "  warnings.warn(\n",
      "/teamspace/studios/this_studio/.venv/lib/python3.11/site-packages/websockets/legacy/__init__.py:6: DeprecationWarning: websockets.legacy is deprecated; see https://websockets.readthedocs.io/en/stable/howto/upgrade.html for upgrade instructions\n",
      "  warnings.warn(  # deprecated in 14.0 - 2024-11-09\n",
      "/teamspace/studios/this_studio/.venv/lib/python3.11/site-packages/uvicorn/protocols/websockets/websockets_impl.py:17: DeprecationWarning: websockets.server.WebSocketServerProtocol is deprecated\n",
      "  from websockets.server import WebSocketServerProtocol\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "* Running on local URL:  http://127.0.0.1:7861\n",
      "* Running on public URL: https://c6b0d9514c5fa7415e.gradio.live\n",
      "\n",
      "This share link expires in 1 week. For free permanent hosting and GPU upgrades, run `gradio deploy` from the terminal in the working directory to deploy to Hugging Face Spaces (https://huggingface.co/spaces)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><iframe src=\"https://c6b0d9514c5fa7415e.gradio.live\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/teamspace/studios/this_studio/.venv/lib/python3.11/site-packages/gradio/routes.py:1341: DeprecationWarning: 'HTTP_422_UNPROCESSABLE_ENTITY' is deprecated. Use 'HTTP_422_UNPROCESSABLE_CONTENT' instead.\n",
      "  return await queue_join_helper(body, request, username)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e8e6068ad35b48fc9853abeea655e01f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/25 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/teamspace/studios/this_studio/.venv/lib/python3.11/site-packages/gradio/routes.py:1341: DeprecationWarning: 'HTTP_422_UNPROCESSABLE_ENTITY' is deprecated. Use 'HTTP_422_UNPROCESSABLE_CONTENT' instead.\n",
      "  return await queue_join_helper(body, request, username)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6c7846e16bbb40dba4011cea3dc47246",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/25 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/teamspace/studios/this_studio/.venv/lib/python3.11/site-packages/gradio/routes.py:1341: DeprecationWarning: 'HTTP_422_UNPROCESSABLE_ENTITY' is deprecated. Use 'HTTP_422_UNPROCESSABLE_CONTENT' instead.\n",
      "  return await queue_join_helper(body, request, username)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2a1a2ce2362343b684828a04026415a6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/25 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import gradio as gr\n",
    "\n",
    "from pruna import PrunaModel\n",
    "\n",
    "# Load PrunaModel\n",
    "pipe = PrunaModel.from_pretrained(\"Sana_600M_512px_diffusers-smashed\")\n",
    "\n",
    "\n",
    "# Inference function\n",
    "def generate_image(prompt):\n",
    "    \"\"\"Generate an image from a given text prompt.\"\"\"\n",
    "    result = pipe(prompt, num_inference_steps=25, guidance_scale=7.5)\n",
    "    return result.images[0]\n",
    "\n",
    "\n",
    "# Create Gradio interface with queueing enabled\n",
    "demo = gr.Interface(\n",
    "    fn=generate_image,\n",
    "    inputs=gr.Textbox(lines=2, placeholder=\"Enter your prompt here...\", label=\"Prompt\"),\n",
    "    outputs=gr.Image(type=\"pil\"),\n",
    "    title=\"Sana Smashed Text-to-Image Demo\",\n",
    "    description=\"Generate high-quality images using the smashed Sana diffusion model optimized with Pruna.\",\n",
    "    allow_flagging=\"never\",\n",
    ")\n",
    "\n",
    "# Enable queueing to handle multiple users\n",
    "demo.queue()\n",
    "\n",
    "# Launch the app\n",
    "if __name__ == \"__main__\":\n",
    "    demo.launch(server_port=7861, share=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusions\n",
    "In this tutorial, we have covered the end-to-end workflow for optimizing and evaluating a text-to-image diffusion model using Pruna.\n",
    "\n",
    "We began by loading the Sana base model and defining the SmashConfig with the desired optimization algorithms and parameters. We then smashed the base model, obtaining an optimized version, and ensured its performance improvements by running an evaluation with the EvaluationAgent.\n",
    "\n",
    "After optimization, we demonstrated how to deploy the smashed model using Gradio to create an interactive demo. This enables anyone to test the model’s capabilities directly in their browser. This end-to-end approach makes it easy to explore trade-offs, iterate on optimization configurations, and deploy robust, production-ready text-to-image models. \n",
    "\n",
    "Check out our other [tutorials](https://docs.pruna.ai/en/stable/docs_pruna/tutorials/index.html) for more examples on optimizing and evaluating large language models, text-to-video models using Pruna."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pruna",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
